package org.ragdollcat.secondaiagent.rag.factory;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
public class MultiQueryExpanderFactory {

    private final ChatClient.Builder builder;

    public MultiQueryExpanderFactory(ChatModel dashscopeChatModel) {
        builder = ChatClient.builder(dashscopeChatModel);
    }


    public List<Query> getMultiQueryExpander(int numberOfQueries, String text) {
        MultiQueryExpander expander = MultiQueryExpander.
                builder().
                chatClientBuilder(builder).
                numberOfQueries(numberOfQueries).
                build();
        return expander.expand(new Query(text));
    }

}
